#include "cpp_co_location_client/IMM_UKF.hpp"

double data1[] = {0.9,0.05,0.05,
				0.05,0.9,0.05,
				0.05,0.05,0.9};
Eigen::Map<Eigen::Matrix<double, 3, 3>> temp_t(data1);
const Eigen::MatrixXd default_interact_pro = temp_t;

Eigen::Vector3d	temp_q(1.0/3.0,1.0/3.0,1.0/3.0);
const Eigen::VectorXd default_model_pro = temp_q;
// const Eigen::MatrixXd default_P = Eigen::MatrixXd::Identity(10,10) * Initial_P_sigma;
double minT = 0.5;
double data2[] = {minT*minT/2,0,0,minT,0,0,1,0,0,0,
    			0,minT*minT/2,0,0,minT,0,0,1,0,0,
    			0,0,minT*minT/2,0,0,minT,0,0,1,minT};
Eigen::Map<Eigen::Matrix<double, 3, 10>> Qtemp(data2);
const Eigen::MatrixXd default_Q = Qtemp.transpose() * Qtemp * 0.25;
const std::vector<std::string> default_model = {"CV","CA","CT"};//

IMM_UKF::IMM_UKF(Eigen::MatrixXd interact_pro, Eigen::VectorXd model_pro, int state_v, std::vector<std::string> model){
	// 设置跟踪ID
	try{
		if(model.size()==0){
			// 如果模型为空，则抛出异常
			throw "valid input !";
		}else{

			// 设置模型数量
			model_size_ = model.size();

			// 遍历每个模型
			for(uint8_t i=0; i<model_size_; ++i){
				// 初始化UKF对象，并加入imm_ukf_向量中
				imm_ukf_.emplace_back(model[i], default_Q, state_v);
				// 打印模型名称（已注释掉）
				//std::cout<<"model ukf "<<model[i]<<std::endl;
			}

			// 设置状态向量维度
			n_x_ = state_v;



			// 初始化合并后的状态向量和协方差矩阵
			x_merge_ = Eigen::VectorXd(n_x_);
			p_merge_ = Eigen::MatrixXd(n_x_,n_x_);

			// 初始化模型概率向量和模型转移概率矩阵
			model_pro_ = model_pro;
			c_ = Eigen::VectorXd(model_size_);

			// // 初始化模型概率为等概率
			// for(int i=0; i< model_size_; ++i){
			// 	model_pro_(i) = 1.0/(float)model_size_;
			// }

			// 初始化模型转移概率矩阵
			interact_pro_ = interact_pro;

			// 初始化每个模型的状态向量和协方差矩阵
			model_X_.resize(model_size_);
			model_P_.resize(model_size_);

			// // 设置每个模型的协方差矩阵为输入参数P
			// for(int i=0; i<model.size(); ++i){
			// 	model_P_[i] = (P);
			// }

			// 初始化合并后的状态向量和协方差矩阵的容器
			X_hat_.resize(model_size_);
			P_hat_.resize(model_size_);

			// 初始化每个模型的状态向量和协方差矩阵
			for(uint8_t i=0; i<model_size_; ++i){
				X_hat_[i] = Eigen::VectorXd(n_x_);
				P_hat_[i] = Eigen::MatrixXd(n_x_,n_x_);
				// 将状态向量和协方差矩阵初始化为0
				X_hat_[i].fill(0.0);
				P_hat_[i].fill(0.0);
			}
		}
	}catch(const char* msg){
		// 输出异常信息
		std::cerr<<msg<<std::endl;
	}
}

bool IMM_UKF::IMM_Initialization(Eigen::Vector3d TargetPosition, uint64_t MeanTime, Eigen::VectorXd model_pro, Eigen::MatrixXd P, Eigen::MatrixXd Q)
{
	// 初始化概率
	model_pro_ = model_pro;
	//初始化UKF的X0，P0 和 时间
	Filtered_Target_State UKF_Output;

	UKF_Output.FilterTime = MeanTime;
	UKF_Output.State = Eigen::VectorXd(n_x_);
	UKF_Output.P = P;

	UKF_Output.State.fill(0.0);
	UKF_Output.State(0) = TargetPosition(0);
	UKF_Output.State(1) = TargetPosition(1);
	UKF_Output.State(2) = TargetPosition(2);

	for(auto &ukf:imm_ukf_)
	{
		ukf.Initialization(UKF_Output, Q);
	}
	isinitialized = true;
	Merged_UKF_Output_ = UKF_Output;
	// 暂时没有失败
	return true;
}

// 处理按时间块分割的量测数据
bool IMM_UKF::Process(Filtered_Target_State &Merged_UKF_Output, std::vector<MeasureMents_In_Period> UKF_Inputs)
{
	std::cout<<"IMM UKF process strat!"<<std::endl;
	Eigen::VectorXd likelihood(model_size_);
	std::vector<Filtered_Target_State> UKF_Outputs(model_size_);
	std::vector<Eigen::VectorXd> States(model_size_);
	std::vector<Eigen::MatrixXd> Covariances(model_size_);

	for(uint8_t i=0;i<model_size_;i++)
	{
		if(!imm_ukf_[i].is_initalized_)
		{
			isinitialized = false;
			return false;
		}
	}

	// 从ukf中取出上次迭代输出的结果 X 和 P，并执行重初始化
	if(State_ReInitialization(States, Covariances,
								interact_pro_, model_pro_))
	{
		std::cout<<"IMM UKF ReInitialization completed!"<<std::endl;
		for(uint8_t i=0;i<model_size_;i++)
		{

			// imm_ukf_[i].Set_P_(Covariances[i]);
			// imm_ukf_[i].Set_X_(States[i]);
			if(!(imm_ukf_[i].Process(UKF_Outputs[i], likelihood(i), 
						UKF_Inputs, States[i], Covariances[i])))
			{
				isinitialized = false;

				std::cout<<"UKF PROCESS failed!:"<<std::endl;
				// imm_ukf_[i].is_initalized_ = false;
				return false;
			}
			std::cout<<"UKF process completed model:"<<int(i+1)<<std::endl;
		}
	}
	else
	{
		for(uint8_t i=0;i<model_size_;i++)
		{
			// imm_ukf_[i].is_initalized_ = false;
		}
		isinitialized = false;
		return false;
	}
	

	if(UpdateMixState(Merged_UKF_Output, model_pro_,
						UKF_Outputs, likelihood))
	{
		if(Merged_UKF_Output.State.hasNaN())
		{
			std::cout<<"UKF failed!:"<<std::endl;
			for(uint8_t i=0;i<model_size_;i++)
			{
				// imm_ukf_[i].is_initalized_ = false;
			}
			isinitialized = false;
			return false;
		}
		Merged_UKF_Output_ = Merged_UKF_Output;
		std::cout<<"Merged_UKF_Output_:"<<Merged_UKF_Output_.State<<std::endl;
		return true;
	}
	else
	{
		std::cout<<"UKF failed!:"<<std::endl;
		for(uint8_t i=0;i<model_size_;i++)
		{
			// imm_ukf_[i].is_initalized_ = false;
		}
		isinitialized = false;
		return false;
	}
}



// 单个模型的重初始化
bool IMM_UKF::State_ReInitialization(std::vector<Eigen::VectorXd> &X_hat, std::vector<Eigen::MatrixXd> &P_hat,
							Eigen::MatrixXd interact_pro, Eigen::VectorXd model_pro)
{
	// 初始化c_数组
	c_.fill(0.0);

	// 遍历每个模型，获取状态和协方差，并计算c_数组
	for(int j=0; j<model_size_; ++j){
		model_X_[j] = imm_ukf_[j].Get_state();
		model_P_[j] = imm_ukf_[j].Get_covariance();
		
		for(int i=0; i<model_size_; ++i){
			c_(j) += interact_pro(i,j)*model_pro(i); 
		}

		// X_hat[j] = imm_ukf_[j].Get_state();
		// P_hat[j] = imm_ukf_[j].Get_covariance();
	}
	// std::cout<<"c_:"<<c_<<std::endl;
	// 遍历每个模型，计算混合状态X_hat和混合协方差P_hat
	for(int j=0; j<model_size_; ++j){
		X_hat[j] = Eigen::VectorXd(n_x_);
		P_hat[j] = Eigen::MatrixXd(n_x_,n_x_);
		X_hat[j].fill(0.0);
		P_hat[j].fill(0.0);
		// std::cout<<"first P_hat:"<<j<<P_hat[j]<<std::endl;
		// std::cout<<"X_hat:"<<X_hat[j]<<std::endl;
		for(int i=0; i<model_size_; ++i){
			double uij =  ((interact_pro(i,j)*model_pro(i))/c_(j));
			// std::cout<<"uij:"<<uij<<std::endl;
			// std::cout<<"IMM UKF Initialized state:"<<model_X_[i]<<std::endl;
			// 根据混合系数u，计算混合状态X_hat
			X_hat[j] += uij * model_X_[i]; 
			// std::cout<<"X_hat:"<<X_hat[j]<<std::endl;
		}
		for(int i=0; i<model_size_; ++i){
			double uij =  (interact_pro(i,j)*model_pro(i))/c_(j);
			// 根据混合系数u，计算混合协方差P_hat
			// std::cout<<"model_P_:"<<i<<model_P_[i]<<std::endl;
			// std::cout<<"X_hat:"<<j<<X_hat[j]<<std::endl;
			// std::cout<<"model_X_:"<<i<<model_X_[i]<<std::endl;
			// std::cout<<"uij:"<<uij<<std::endl;
			Eigen::MatrixXd temp = (model_X_[i] - X_hat[j])*(model_X_[i] - X_hat[j]).transpose();
			// std::cout<<"temp"<<temp<<std::endl;
			P_hat[j] += (uij * (model_P_[i] + temp));
			// std::cout<<"P_hat:"<<j<<P_hat[j]<<std::endl;
		}
		
	}



	
	return true;
}


bool IMM_UKF::UpdateMixState(Filtered_Target_State &Merged_UKF_Output, Eigen::VectorXd &model_pro,
					std::vector<Filtered_Target_State> UKF_Outputs, Eigen::VectorXd model_ita)				
{

	// 初始化临时变量c_temp
	float c_temp = 0;
	auto model_p = Eigen::VectorXd(model_size_);
	model_p.fill(0.0);
	// 计算混合概率的权重和
	for(int i=0; i<model_size_; ++i){
		c_temp += model_ita[i] * c_(i);
	}
	// std::cout<<"model_ita:"<<model_ita<<std::endl;
	// std::cout<<"c_:"<<c_<<std::endl;
	// 计算每个模型的混合概率
	for(int i=0; i<model_size_; ++i){
		model_p(i) = (model_ita[i]*c_(i))/c_temp;
	}
	model_pro = model_p;

	std::cout<<"model_pro:"<<model_pro<<std::endl;
	// 检查模型概率是否合法
	if(model_pro.hasNaN())
	{	
		// std::abort();
		return false;
	}

	// 归一化模型概率
	if(abs(model_pro.sum()) !=0)
		model_pro /=model_pro.sum();
	else
	{
		return false;
	}

	// 初始化合并后的状态向量和协方差矩阵为0
	Merged_UKF_Output.State = Eigen::VectorXd(n_x_);
	Merged_UKF_Output.P = Eigen::MatrixXd(n_x_,n_x_);
	Merged_UKF_Output.State.fill(0.0);
	Merged_UKF_Output.P.fill(0.0);
	Merged_UKF_Output.FilterTime = uint64_t(0);

	// 遍历每个模型
	for(int i=0; i<model_size_; ++i){
		// 获取当前模型的状态向量和协方差矩阵
		model_X_[i] = UKF_Outputs[i].State;
		model_P_[i] = UKF_Outputs[i].P;

		// 根据模型概率加权计算合并后的状态向量
		Merged_UKF_Output.State += model_X_[i] * model_pro_(i);
		Merged_UKF_Output.FilterTime+=UKF_Outputs[i].FilterTime;
	}

	Merged_UKF_Output.FilterTime /= uint64_t(model_size_);

	// 再次遍历每个模型
	for(int i=0; i<model_size_; ++i){
		// 根据模型概率和状态偏差计算合并后的协方差矩阵
		Merged_UKF_Output.P += model_pro_(i) * (model_P_[i] + (model_X_[i] -Merged_UKF_Output.State)* (model_X_[i] -Merged_UKF_Output.State).transpose());
	}

	return true;
}



































